﻿
#include "Application.h"
#include "NavierStokesEquation.h"
#include "Mesh.h"
#include "Timer.h"
#include "Transient.h"

#include "Algorithm.h"


VTKM_EXEC
vtkm::Vec<FloatDefault, 5> valueExactCoutteFlow(FloatDefault t, const Vec3f& p) 
{
  FloatDefault _gamma = 1.4, _prandtl = 0.72, _mach = 0.2;

  FloatDefault x = p[0];
  FloatDefault y = p[1];
  FloatDefault z = p[2];

  FloatDefault T2 = 0.85, T1 = 0.8;
  FloatDefault tem =
    T1 + (T2 - T1) * y / 2 + 0.5 * _prandtl * (_gamma - 1) * _mach * _mach * y / 2 * (1 - y / 2);
  FloatDefault pre = 1. / (_gamma * _mach * _mach);
  FloatDefault density = pre * _gamma * _mach * _mach / tem;
  FloatDefault u = y / 2.0, v = 0, w = 0;
  Vec3f mom{ density * u, density * v, density * w };
  FloatDefault rhoe = pre / (_gamma - 1) + 0.5 * density * (u * u + v * v + w * w);

  vtkm::Vec<FloatDefault, 5> value;
  value[0] = density;
  value[1] = density * u;
  value[2] = density * v;
  value[3] = density * w;
  value[4] = rhoe;

  return value;
}

struct CoutteFlowIC : WorkletInitCondition<CoutteFlowIC>
{
  VTKM_EXEC ValueType ComputeValue(const Vec3f& p) const
  {
    FloatDefault gam = 1.4, gamm1 = gam - 1, mach = 0.2;
    FloatDefault u = 1;
    FloatDefault v = 0;
    FloatDefault density = 1;
    FloatDefault pre = 1. / (gam * mach * mach);

    ValueType value;
    value[0] = density;
    value[1] = density * u;
    value[2] = density * v;
    value[3] = density * 0;
    value[4] = pre / gamm1 + 0.5 * density * (u * u + v * v);

    return value;
  }
};




struct ComputeResidual_Coutte : vtkm::worklet::WorkletCellNeighborhood
{
  VTKM_CONT
  ComputeResidual_Coutte(FloatDefault t)
    : _t(t)
  {
  }

  using ControlSignature =
    void(CellSetIn, FieldInNeighborhood, FieldOut Rhs, FieldIn centroid, FieldIn spacing);
  using ExecutionSignature = void(Boundary, _2, _3, _4, _5);

  template<typename NeighIn, typename RhsType>
  VTKM_EXEC void operator()(const vtkm::exec::BoundaryState& boundary,
                            const NeighIn& value,
                            RhsType& rhs,
                            const Vec3f& centroid,
                            const Vec3f& spacing) const
  {
    FloatDefault penalty = 4;
    Vec3f normal{ 0, 0, 0 };
    DataPack<RhsType::NUM_COMPONENTS> dl, dr, lift;
    {
      //auto& ul = value.Get(0, 0, 0);
      //auto ur =  value.Get(-1, 0, 0);
      dl.uh = value.Get(0, 0, 0);
      dr.uh = value.Get(-1, 0, 0);
      RhsType lflux, rflux;
      if (boundary.IJK[0] == 0)
      {
        Vec3f p = centroid;
        p[0] = p[0] - 0.5 * spacing[0];
        //ur = ul;
        //dr.uh = value.Get(0, 0, 0);
        dr.uh = valueExactCoutteFlow(_t, p);
      }
      dl.InvisTerm();
      dr.InvisTerm();

      normal = Vec3f{ -1, 0, 0 };
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
      {
        lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
        lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
      }

      RiemannFlux(dl, dr, normal, lflux);
      lift.ViscousTerm();
      //std::cout << lift.vis_term[0] << std::endl;
      //std::cout << lift.vis_term[1] << std::endl;
      //std::cout << lift.vis_term[2] << std::endl;
      //std::cout << lift.vis_term[3] << std::endl;
      //std::cout << lift.vis_term[4] << std::endl;

      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
      {
        lflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
      }


      dl.uh = value.Get(0, 0, 0);
      dr.uh = value.Get(1, 0, 0);
      if (boundary.IJK[0] == boundary.PointDimensions[0] - 1)
      {
        Vec3f p = centroid;
        p[0] = p[0] + 0.5 * spacing[0];
        //ur = ul;
        //dr.uh = value.Get(0, 0, 0);
        dr.uh = valueExactCoutteFlow(_t, p);
      }
      dl.InvisTerm();
      dr.InvisTerm();


      normal = Vec3f{ 1, 0, 0 };
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
      {
        lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
        lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
      }

      RiemannFlux(dl, dr, normal, rflux);
      lift.ViscousTerm();
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
      {
        rflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
      }

      //RiemannFlux(ul, ur, Vec3f{ 1, 0, 0 }, rflux);
      rhs = (rflux + lflux) / spacing[0];
    }

    {
      //auto& ul = value.Get(0, 0, 0);
      //auto ur = value.Get(0, -1, 0);
      dl.uh = value.Get(0, 0, 0);
      dr.uh = value.Get(0, -1, 0);
      RhsType lflux, rflux;
      if (boundary.IJK[1] == 0)
      {
        Vec3f p = centroid;
        p[1] = p[1] - 0.5 * spacing[1];
        // dr.uh = value.Get(0, 0, 0);
        dr.uh = valueExactCoutteFlow(_t, p);
      }

      //dl.uh = ul;
      //dr.uh = ur;
      dl.InvisTerm();
      dr.InvisTerm();

      normal = Vec3f{ 0, -1, 0 };
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
      {
        lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
        lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
      }

      RiemannFlux(dl, dr, normal, lflux);
      lift.ViscousTerm();
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
      {
        lflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
      }

      //ur = value.Get(0, 1, 0);
      dl.uh = value.Get(0, 0, 0);
      dr.uh = value.Get(0, 1, 0);
      if (boundary.IJK[1] == boundary.PointDimensions[1] - 1)
      {
        Vec3f p = centroid;
        p[1] = p[1] + 0.5 * spacing[1];
        //dr.uh = value.Get(0, 0, 0);
        dr.uh = valueExactCoutteFlow(_t, p);
      }
      dl.InvisTerm();
      dr.InvisTerm();

      normal = Vec3f{ 0, 1, 0 };
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
      {
        lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
        lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
      }

      RiemannFlux(dl, dr, normal, rflux);
      lift.ViscousTerm();
      for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
      {
        rflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
      }

      rhs += (rflux + lflux) / spacing[1];
    }
  }


  
  template<typename DataPackType, typename VType>
  VTKM_EXEC void RiemannFlux(const DataPackType& dl,
                             const DataPackType& dr,
                             const Vec3f& normal,
                             VType& riemann_flux) const
  {
    for (IdComponent eq = 0; eq < VType::NUM_COMPONENTS; eq++)
    {
      riemann_flux[eq] =
        vtkm::Dot(dl.inv_term[eq] + dr.inv_term[eq], normal) / 2 + 0.6 * (dl.uh[eq] - dr.uh[eq]);
    }
  }
  FloatDefault _t;
};


RegisterObject(NavierStokesEquation);
NavierStokesEquation::NavierStokesEquation(const Configuration& cfg)
  : EulerEquation(cfg)
{
}

void NavierStokesEquation::Init()
{
  EulerEquation::Init();
  /*vtkm::cont::Invoker invoker;
  const auto& centroid = _mesh.GetCellField("center");

  vtkm::cont::Invoker{}(MeshModifier_ForwardStep{}, _mesh.GetCellSet(), centroid, _blk_id, _bnd_id);*/
}

void NavierStokesEquation::InitialCondition()
{
  vtkm::cont::Timer timer;
  timer.Start();
  vtkm::cont::Invoker invoker;
  invoker(CoutteFlowIC{}, _mesh.GetCellSet(), _solution, _mesh.GetCoordinate());

  console::Success("NavierStokesEquation初始条件设置完成", timer.GetElapsedTime(), "s");
}



void NavierStokesEquation::UpdateResidual()
{
  auto& transient = static_cast<Transient&>(*(_app.GetExecutioner()));
  Real current_time = transient.CurrentTime();

  Timer timer;
  timer.Start();
  const auto& centroid = _mesh.GetCellField("center");
  const auto& spacing = _mesh.GetCellField("spacing");
  vtkm::cont::Invoker{}(
    ComputeResidual_Coutte{ current_time }, _mesh.GetCellSet(), _solution, _rhs, centroid, spacing);
  console::Info("EulerianSystem::UpdateResidual运行完成，", "时间: ", timer.GetElapsedTime(), " s");
}
